import streamlit as st

from util.hunyuan_chat import chat
from util.hunyuan_chat import Msg

str = ""
async def handle(x):
    global str
    async for xx in x:
        str += xx
        yield xx


if "chat" not in st.session_state:
    st.session_state["chat"] = []


messages = st.container()
if prompt := st.chat_input("Say something"):
    st.session_state.chat.append(Msg("user", prompt))
    for msg in st.session_state.chat:
        messages.chat_message(msg.role).write(msg.content)
    messages.chat_message("assistant").write_stream(handle(chat(st.session_state.chat)))
    st.session_state.chat.append(Msg("assistant", str))